import numpy as np

from calc_force import *
from read_data import *
from write_data import *
from plot_data import *
from filter_data import *
from transform_coord import *
from check_exist import *
from time_cut import *

# 是否保存
save_flag = False

# F16/F22/F22_STATIC/F22_SWEEP/FlyWing/
MODEL = "F22_STATIC"  

DATASET_PREFIX = "../DATASET_ALL/DATASET_UNCUT/"
AERO_DIR = DATASET_PREFIX + MODEL + "/AERO/"
ANGLE_DIR = DATASET_PREFIX + MODEL + "/ANGLE/"
TENSION_DIR = DATASET_PREFIX + MODEL + "/TENSION/"
ENCODER_DIR = DATASET_PREFIX + MODEL + "/ENCODER/"

# 读取文件列表
print("=========项目读取=========")
# file_list = get_l_folder(INPUT_DIR)
file_list = get_l_list(MODEL)
# file_list = ["A04","A07"]


# 获取频率信息
if(MODEL=="F22"):
    freq_motion_dict = get_f(MODEL, DATASET_PREFIX)
freq_aero = 1
freq_tension = 10
freq_encoder = 20
freq_angle = 200

# 气动力计算迭代轮数
iter_num = 5


# 检查文件完整性
print("========完整性检查========")
check_e(file_list, MODEL, DATASET_PREFIX)


for file_name in file_list:
    print("==========="+file_name+"===========")

    # 读取气动数据（电压值）
    u_nw_raw = read_d_aero(AERO_DIR + file_name + "_NW.txt")
    u_w_raw = read_d_aero(AERO_DIR + file_name + "_W.txt")

    # 读取拉力数据
    tension_nw = []
    tension_w =[]

    # 读取编码器数据
    encoder_nw = []
    encoder_w = []

    # 读取姿态角数据
    angle_nw = read_d_angle(ANGLE_DIR + file_name + "_NW*.csv")
    angle_w = read_d_angle(ANGLE_DIR + file_name + "_W*.csv")

    angle_nw = filter_angle(angle_nw)
    angle_w = filter_angle(angle_w)

    # 时间对齐
    freq_angle = 200
    
    # 不滤波
    u_nw = u_nw_raw
    u_w = u_w_raw

    F_n_Q_aero = []
    F_n_Q_body = []
    F_n_Q_wind = []

    # 计算气动力(机体轴)
    for i in range(len(u_nw)):
        
        u0 = u_nw[i]
        u1 = u_w[i]

        # X, Y, Z, Mx, My, Mz, Q
        f,q = calc_f2(u0, u1, iter_num)

        f = np.append(f,q)
        F_n_Q_aero.append(f)

    F_n_Q_aero = np.array(F_n_Q_aero)

    # 坐标系转换(天平轴→体轴系)
    F_n_Q_body = transform_a2b(F_n_Q_aero)

    # 坐标系转换(体轴系→气流系)
    Angle_body = angle_w
    F_n_Q_wind, Angle_wind = transform_b2w(F_n_Q_body,Angle_body)

    # 计算气动力系数
    C_n_Q_body = calc_coeff(F_n_Q_body, MODEL)
    C_n_Q_wind = calc_coeff(F_n_Q_wind, MODEL)

    if(save_flag==False):
        
        # 画图
        # plot_all(angle_nw,angle_w,freq_angle,F_n_Q_body,F_n_Q_wind,freq_aero,freq_motion,angle_w)
        # plot_all_time(angle_nw,angle_w,freq_angle,C_n_Q_body,C_n_Q_wind,freq_aero)
        # plot_all_time(angle_nw,angle_w,freq_angle,F_n_Q_aero,F_n_Q_aero,freq_aero)


        # 时间-电压
        # plot_u(u_nw_raw, u_w_raw, u_nw, u_w, freq_aero)

        # # 时间-角度
        plot_comp(angle_nw,angle_w,freq_angle)

        # # 时间-气动力
        # plot_f_time(F_n_Q_body, freq_aero)
        # plot_f_time(F_n_Q_wind, freq_aero)

        plot_f_time(C_n_Q_body, freq_aero)
        plot_f_time(C_n_Q_wind, freq_aero)

        # 角度-气动力
        # plot_all_angle(angle_nw,angle_w,freq_angle,-C_n_Q_body,-C_n_Q_wind,freq_aero,freq_motion,angle_w)
        # plot_f_degree_ani(F_n_Q_body, angle_w, freq_motion, freq_angle, False)
        # plot_f_degree(C_n_Q_body, angle_w, freq_motion, freq_angle, False)

        # 气动角-气动力
        # plot_f_degree_wind(F_n_Q_wind, Angle_wind, freq_motion, freq_angle, False)
        # plot_f_degree_wind_all(-C_n_Q_wind, Angle_wind, freq_motion, freq_angle, True)

        # plot_all_angle_ani(angle_nw,angle_w,freq_angle,F_n_Q_body,F_n_Q_wind,freq_aero,freq_motion,angle_w)

    else:
        # 保存数据到文件
        force_list = [F_n_Q_aero, F_n_Q_body, F_n_Q_wind]
        coeff_list = [C_n_Q_body, C_n_Q_wind]
        anlge_list = [Angle_body, Angle_wind]
        pic_list = []
        all_list = [force_list, coeff_list, anlge_list, pic_list]
        write_all(all_list, MODEL,file_name)
    
    

